package v1

import (
	"crypto/hmac"
	"crypto/sha1"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

const (
	// StsSignVersion sts sign version
	StsSignVersion = "1.0"
	// StsAPIVersion sts api version
	StsAPIVersion = "2015-04-01"
	// StsHost sts host, detail information:https://help.aliyun.com/document_detail/66053.html?spm=a2c4g.11186623.2.12.151338dfkKJJKh
	StsHost = "https://sts.aliyuncs.com/"
	// TimeFormat time fomrat
	TimeFormat = "2006-01-02T15:04:05Z"
	// RespBodyFormat  respone body format
	RespBodyFormat = "JSON"
	// PercentEncode '/'
	PercentEncode = "%2F"
	// HTTPGet http get method
	HTTPGet = "GET"
)

// token server config
type StsConfig struct {
	AccessKeyId     string `json:"AccessKeyID"`
	AccessKeySecret string `json:"AccessKeySecret"`
	RoleArn         string `json:"RoleArn"`
	TokenExpireTime string `json:"TokenExpireTime"`
	PolicyFile      string `json:"PolicyFile"`
}

var stsConfig StsConfig

// Policy allows not to be configured
var stsPolicy string

// you can modify the value according to the actual situation
var stsSession string = "test_token_server"

// StsTokenInfo for app client get success
type StsTokenInfo struct {
	StatusCode      int    `json:"StatusCode"`
	AccessKeyId     string `json:"AccessKeyId"`
	AccessKeySecret string `json:"AccessKeySecret"`
	SecurityToken   string `json:"SecurityToken"`
	Expiration      string `json:"Expiration"`
}

// StsErrorInfo for app client get failure
type StsErrorInfo struct {
	StatusCode   int    `json:"StatusCode"`
	ErrorCode    string `json:"ErrorCode"`
	ErrorMessage string `json:"ErrorMessage"`
}

// Credentials for token server get success
type Credentials struct {
	AccessKeyId     string `json:"AccessKeyId"`
	AccessKeySecret string `json:"AccessKeySecret"`
	Expiration      string `json:"Expiration"`
	SecurityToken   string `json:"SecurityToken"`
}

// AssumedRoleUser for token server get success
type AssumedRoleUser struct {
	Arn           string `json:"Arn"`
	AssumedRoleId string `json:"AssumedRoleId"`
}

// StsResponse the response of sts service
type StsResponse struct {
	Credentials     Credentials     `json:"Credentials"`
	AssumedRoleUser AssumedRoleUser `json:"AssumedRoleUser"`
	RequestId       string          `json:"RequestId"`
}

func randomInt(min, max int) int {
	return min + rand.Intn(max-min)
}

func randomString(len int) string {
	bytes := make([]byte, len)
	for i := 0; i < len; i++ {
		bytes[i] = byte(randomInt(65, 90))
	}
	return string(bytes)
}

// Private function
func generateSignedURL(sessionName string, policy string) (string, error) {

	rand.Seed(time.Now().UnixNano())
	queryStr := "SignatureVersion=" + StsSignVersion
	queryStr += "&Format=" + RespBodyFormat
	queryStr += "&Timestamp=" + url.QueryEscape(time.Now().UTC().Format(TimeFormat))
	queryStr += "&RoleArn=" + url.QueryEscape(stsConfig.RoleArn)
	queryStr += "&RoleSessionName=" + sessionName
	queryStr += "&AccessKeyId=" + stsConfig.AccessKeyId
	queryStr += "&SignatureMethod=HMAC-SHA1"
	queryStr += "&Version=" + StsAPIVersion
	queryStr += "&Action=AssumeRole"
	queryStr += "&SignatureNonce=" + randomString(10)
	queryStr += "&DurationSeconds=" + stsConfig.TokenExpireTime
	if policy != "" {
		queryStr += "&Policy=" + url.QueryEscape(policy)
	}

	// Sort query string
	queryParams, err := url.ParseQuery(queryStr)
	if err != nil {
		return "", err
	}
	sortUrl := strings.Replace(queryParams.Encode(), "+", "%20", -1)
	strToSign := HTTPGet + "&" + PercentEncode + "&" + url.QueryEscape(sortUrl)

	// Generate signature
	hashSign := hmac.New(sha1.New, []byte(stsConfig.AccessKeySecret+"&"))
	hashSign.Write([]byte(strToSign))
	signature := base64.StdEncoding.EncodeToString(hashSign.Sum(nil))

	// Build url
	assumeURL := StsHost + "?" + queryStr + "&Signature=" + url.QueryEscape(signature)
	return assumeURL, nil
}

func sendRequest(url string) ([]byte, int, error) {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	client := &http.Client{Transport: tr}

	resp, err := client.Get(url)
	if err != nil {
		return nil, -1, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	return body, resp.StatusCode, err
}

func GetStsToken(w http.ResponseWriter, r *http.Request) {
	url, err := generateSignedURL(stsSession, stsPolicy)
	stsErrorInfo := &StsErrorInfo{}
	if err != nil {
		stsErrorInfo.StatusCode = http.StatusInternalServerError
		stsErrorInfo.ErrorCode = "internal server error"
		stsErrorInfo.ErrorMessage = err.Error()
		bs, _ := json.Marshal(stsErrorInfo)
		http.Error(w, string(bs), http.StatusInternalServerError)
		return
	}

	body, status, err := sendRequest(url)
	if err != nil {
		stsErrorInfo.StatusCode = http.StatusInternalServerError
		stsErrorInfo.ErrorCode = "internal server error"
		stsErrorInfo.ErrorMessage = err.Error()
		bs, _ := json.Marshal(stsErrorInfo)
		http.Error(w, string(bs), http.StatusInternalServerError)
		return
	}

	if status != http.StatusOK {
		stsErrorInfo := &StsErrorInfo{}
		stsErrorInfo.StatusCode = status
		stsErrorInfo.ErrorCode = "get sts token failure"
		stsErrorInfo.ErrorMessage = string(body)
		bs, _ := json.Marshal(stsErrorInfo)
		http.Error(w, string(bs), status)
		return
	}

	resp := StsResponse{}
	err = json.Unmarshal(body, &resp)
	if err != nil {
		stsErrorInfo := &StsErrorInfo{}
		stsErrorInfo.StatusCode = http.StatusInternalServerError
		stsErrorInfo.ErrorCode = "internal server error"
		stsErrorInfo.ErrorMessage = err.Error()
		bs, _ := json.Marshal(stsErrorInfo)
		http.Error(w, string(bs), http.StatusInternalServerError)
		return
	}

	stsTokenInfo := &StsTokenInfo{}
	stsTokenInfo.StatusCode = http.StatusOK
	stsTokenInfo.AccessKeyId = resp.Credentials.AccessKeyId
	stsTokenInfo.AccessKeySecret = resp.Credentials.AccessKeySecret
	stsTokenInfo.SecurityToken = resp.Credentials.SecurityToken
	stsTokenInfo.Expiration = resp.Credentials.Expiration
	bs, _ := json.Marshal(stsTokenInfo)
	w.Write(bs)
	return
}

func main() {
	if len(os.Args) < 4 {
		fmt.Printf("usage:%s ip port config\n", os.Args[0])
		return
	}
	bindHost := os.Args[1] + ":" + os.Args[2]

	// read config file
	data, err := ioutil.ReadFile(os.Args[3])
	if err != nil {
		fmt.Printf("read %s error,%s\n", os.Args[3], err.Error())
		return
	}

	// parser config
	err = json.Unmarshal(data, &stsConfig)
	if err != nil {
		fmt.Printf(" json.Unmarshal error,%s\n", err.Error())
		return
	}

	// read policy
	if stsConfig.PolicyFile != "" {
		data, err = ioutil.ReadFile(stsConfig.PolicyFile)
		if err != nil {
			fmt.Printf("ioutil.ReadFile %s error,%s\n", stsConfig.PolicyFile, err.Error())
			return
		}
		stsPolicy = string(data)
	}
	stsPolicy = strings.Replace(stsPolicy, "\r", "", -1)
	stsPolicy = strings.Replace(stsPolicy, "\n", "", -1)

	//http.HandleFunc("/", getStsToken)

	// attention: you must use http.ListenAndServeTLS to start https server for security
	// start http server for sample demo
	err = http.ListenAndServe(bindHost, nil)
	if err != nil {
		fmt.Println("ListenAndServe error:", err)
	}
}
